import tensorflow as tf
import os

if __name__ == '__main__':
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    n = 6
    input = tf.random.normal(shape=[n, 28 * 28])

    W1 = tf.Variable(tf.random.truncated_normal(shape=[784, 256]))
    b1 = tf.Variable(tf.zeros(256))

    W2 = tf.Variable(tf.random.truncated_normal(shape=[256, 128]))
    b2 = tf.Variable(tf.zeros(128))

    W3 = tf.Variable(tf.random.truncated_normal(shape=[128, 64]))
    b3 = tf.Variable(tf.zeros(64))

    W4 = tf.Variable(tf.random.truncated_normal(shape=[64, 10]))
    b4 = tf.Variable(tf.zeros(10))

    hidden1 = tf.nn.relu(input @ W1 + b1)
    hidden2 = tf.nn.relu(hidden1 @ W2 + b2)
    hidden3 = tf.nn.relu(hidden2 @ W3 + b3)
    output = tf.nn.relu(hidden3 @ W4 + b4)

    print(output)